{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e34cac49",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# This implementation is based on the DenseNet-BC implementation in torchvision\n",
    "# https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py\n",
    "\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.checkpoint as cp\n",
    "from collections import OrderedDict\n",
    "\n",
    "\n",
    "def _bn_function_factory(norm, relu, conv):\n",
    "    def bn_function(*inputs):\n",
    "        concated_features = torch.cat(inputs, 1)\n",
    "        bottleneck_output = conv(relu(norm(concated_features)))\n",
    "        return bottleneck_output\n",
    "\n",
    "    return bn_function\n",
    "\n",
    "\n",
    "class _DenseLayer(nn.Module):\n",
    "    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, efficient=False):\n",
    "        super(_DenseLayer, self).__init__()\n",
    "        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),\n",
    "        self.add_module('relu1', nn.ReLU(inplace=True)),\n",
    "        self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * growth_rate,\n",
    "                        kernel_size=1, stride=1, bias=False)),\n",
    "        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),\n",
    "        self.add_module('relu2', nn.ReLU(inplace=True)),\n",
    "        self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,\n",
    "                        kernel_size=3, stride=1, padding=1, bias=False)),\n",
    "        self.drop_rate = drop_rate\n",
    "        self.efficient = efficient\n",
    "\n",
    "    def forward(self, *prev_features):\n",
    "        bn_function = _bn_function_factory(self.norm1, self.relu1, self.conv1)\n",
    "        if self.efficient and any(prev_feature.requires_grad for prev_feature in prev_features):\n",
    "            bottleneck_output = cp.checkpoint(bn_function, *prev_features)\n",
    "        else:\n",
    "            bottleneck_output = bn_function(*prev_features)\n",
    "        new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))\n",
    "        if self.drop_rate > 0:\n",
    "            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)\n",
    "        return new_features\n",
    "\n",
    "\n",
    "class _Transition(nn.Sequential):\n",
    "    def __init__(self, num_input_features, num_output_features):\n",
    "        super(_Transition, self).__init__()\n",
    "        self.add_module('norm', nn.BatchNorm2d(num_input_features))\n",
    "        self.add_module('relu', nn.ReLU(inplace=True))\n",
    "        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,\n",
    "                                          kernel_size=1, stride=1, bias=False))\n",
    "        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))\n",
    "\n",
    "\n",
    "class _DenseBlock(nn.Module):\n",
    "    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, efficient=False):\n",
    "        super(_DenseBlock, self).__init__()\n",
    "        for i in range(num_layers):\n",
    "            layer = _DenseLayer(\n",
    "                num_input_features + i * growth_rate,\n",
    "                growth_rate=growth_rate,\n",
    "                bn_size=bn_size,\n",
    "                drop_rate=drop_rate,\n",
    "                efficient=efficient,\n",
    "            )\n",
    "            self.add_module('denselayer%d' % (i + 1), layer)\n",
    "\n",
    "    def forward(self, init_features):\n",
    "        features = [init_features]\n",
    "        for name, layer in self.named_children():\n",
    "            new_features = layer(*features)\n",
    "            features.append(new_features)\n",
    "        return torch.cat(features, 1)\n",
    "\n",
    "\n",
    "class DenseNet(nn.Module):\n",
    "    r\"\"\"Densenet-BC model class, based on\n",
    "    `\"Densely Connected Convolutional Networks\" <https://arxiv.org/pdf/1608.06993.pdf>`\n",
    "    Args:\n",
    "        growth_rate (int) - how many filters to add each layer (`k` in paper)\n",
    "        block_config (list of 3 or 4 ints) - how many layers in each pooling block\n",
    "        num_init_features (int) - the number of filters to learn in the first convolution layer\n",
    "        bn_size (int) - multiplicative factor for number of bottle neck layers\n",
    "            (i.e. bn_size * k features in the bottleneck layer)\n",
    "        drop_rate (float) - dropout rate after each dense layer\n",
    "        num_classes (int) - number of classification classes\n",
    "        small_inputs (bool) - set to True if images are 32x32. Otherwise assumes images are larger.\n",
    "        efficient (bool) - set to True to use checkpointing. Much more memory efficient, but slower.\n",
    "    \"\"\"\n",
    "    def __init__(self, growth_rate=12, block_config=(16, 16, 16), compression=0.5,\n",
    "                 num_init_features=24, bn_size=4, drop_rate=0,\n",
    "                 num_classes=10, small_inputs=True, efficient=False):\n",
    "\n",
    "        super(DenseNet, self).__init__()\n",
    "        assert 0 < compression <= 1, 'compression of densenet should be between 0 and 1'\n",
    "\n",
    "        # First convolution\n",
    "        if small_inputs:\n",
    "            self.features = nn.Sequential(OrderedDict([\n",
    "                ('conv0', nn.Conv2d(3, num_init_features, kernel_size=3, stride=1, padding=1, bias=False)),\n",
    "            ]))\n",
    "        else:\n",
    "            self.features = nn.Sequential(OrderedDict([\n",
    "                ('conv0', nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)),\n",
    "            ]))\n",
    "            self.features.add_module('norm0', nn.BatchNorm2d(num_init_features))\n",
    "            self.features.add_module('relu0', nn.ReLU(inplace=True))\n",
    "            self.features.add_module('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1,\n",
    "                                                           ceil_mode=False))\n",
    "\n",
    "        # Each denseblock\n",
    "        num_features = num_init_features\n",
    "        for i, num_layers in enumerate(block_config):\n",
    "            block = _DenseBlock(\n",
    "                num_layers=num_layers,\n",
    "                num_input_features=num_features,\n",
    "                bn_size=bn_size,\n",
    "                growth_rate=growth_rate,\n",
    "                drop_rate=drop_rate,\n",
    "                efficient=efficient,\n",
    "            )\n",
    "            self.features.add_module('denseblock%d' % (i + 1), block)\n",
    "            num_features = num_features + num_layers * growth_rate\n",
    "            if i != len(block_config) - 1:\n",
    "                trans = _Transition(num_input_features=num_features,\n",
    "                                    num_output_features=int(num_features * compression))\n",
    "                self.features.add_module('transition%d' % (i + 1), trans)\n",
    "                num_features = int(num_features * compression)\n",
    "\n",
    "        # Final batch norm\n",
    "        self.features.add_module('norm_final', nn.BatchNorm2d(num_features))\n",
    "\n",
    "        # Linear layer\n",
    "        self.classifier = nn.Linear(num_features, num_classes)\n",
    "\n",
    "        # Initialization\n",
    "        for name, param in self.named_parameters():\n",
    "            if 'conv' in name and 'weight' in name:\n",
    "                n = param.size(0) * param.size(2) * param.size(3)\n",
    "                param.data.normal_().mul_(math.sqrt(2. / n))\n",
    "            elif 'norm' in name and 'weight' in name:\n",
    "                param.data.fill_(1)\n",
    "            elif 'norm' in name and 'bias' in name:\n",
    "                param.data.fill_(0)\n",
    "            elif 'classifier' in name and 'bias' in name:\n",
    "                param.data.fill_(0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        features = self.features(x)\n",
    "        out = F.relu(features, inplace=True)\n",
    "        out = F.adaptive_avg_pool2d(out, (1, 1))\n",
    "        out = torch.flatten(out, 1)\n",
    "        out = self.classifier(out)\n",
    "        return out"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}